import os
import logging
import shutil

import numpy as np
import torch
import torch.backends.cudnn as cudnn
from PIL import Image
import matplotlib.pyplot as plt
from colorama import Fore, Style  # 用于颜色标记


class SiamesePredictor:
    """
    Siamese 网络模型类，用于加载模型并进行图像相似度检测。

    Attributes:
        input_shape (list): 输入图像的尺寸，默认为 [105, 105]。
        cuda (bool): 是否使用 CUDA，默认为 True。
        model_path (str): 模型权重路径。
        model (torch.nn.Module): Siamese 网络模型实例。
        save_result_img_dir (str): 保存结果的目录。
    """

    _defaults = {
        "input_shape": [105, 105],  # 输入图像尺寸
        "cuda": True,  # 是否使用 CUDA
        "background_color": (128, 128, 128),  # 背景颜色
    }

    @classmethod
    def get_defaults(cls, n):
        """获取默认配置项的值"""
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            raise AttributeError(f"Unrecognized attribute name '{n}'")

    def __init__(self, model_path, match_model, save_result_img_name, **kwargs):
        """
        初始化 Siamese 模型。

        Args:
            model_path (str): 模型权重路径。
            match_model (torch.nn.Module): 模型实例。
            save_result_img_name (str): 保存结果的文件名。
            **kwargs: 可覆盖默认配置的参数。
        """
        self.model_path = model_path  # 从外界传入模型路径
        self.model = match_model
        self.save_result_img_dir = save_result_img_name
        self.__dict__.update(self._defaults)
        for name, value in kwargs.items():
            setattr(self, name, value)
        self.generate()

    def generate(self):
        """加载模型权重并初始化模型"""
        logging.info("Loading weights into state dict...")
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        try:
            self.model.load_state_dict(torch.load(self.model_path, map_location=device), strict=False)
            self.net = self.model.eval()
            logging.info(f"{self.model_path} model loaded.")
        except Exception as e:
            logging.error(f"Failed to load model: {e}")
            raise

        if self.cuda:
            self.net = torch.nn.DataParallel(self.net)
            cudnn.benchmark = True
            self.net = self.net.cuda()

    def letterbox_image(self, image, size):
        """对图像进行不失真 resize。

        Args:
            image (PIL.Image): 输入的图像。
            size (tuple): 目标尺寸 (width, height)。

        Returns:
            PIL.Image: 处理后的图像。
        """
        image = image.convert("RGB")
        iw, ih = image.size
        w, h = size
        scale = min(w / iw, h / ih)
        nw, nh = int(iw * scale), int(ih * scale)

        image = image.resize((nw, nh), Image.BICUBIC)
        new_image = Image.new("RGB", size, self.background_color)
        new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))
        if self.input_shape[-1] == 1:
            new_image = new_image.convert("L")
        return new_image

    def preprocess_image(self, image):
        """预处理图像，包括 resize 和归一化。

        Args:
            image (PIL.Image): 输入的图像。

        Returns:
            torch.Tensor: 处理后的图像张量。
        """
        image = self.letterbox_image(image, [self.input_shape[1], self.input_shape[0]])
        image = self.preprocess_input(np.array(image, np.float32))
        image = torch.from_numpy(np.expand_dims(np.transpose(image, (2, 0, 1)), 0)).type(torch.FloatTensor)
        if self.cuda:
            image = image.cuda()
        return image

    def detect_image(self, image_1: Image.Image, image_2: Image.Image) -> float:
        """检测两张图片的相似度。

        Args:
            image_1 (PIL.Image): 第一张图片。
            image_2 (PIL.Image): 第二张图片。

        Returns:
            float: 相似度值。
        """
        # 将图像转换为 RGB 格式
        image_1 = self.convert_to_rgb(image_1)
        image_2 = self.convert_to_rgb(image_2)

        # 预处理图像
        photo_1 = self.preprocess_image(image_1)
        photo_2 = self.preprocess_image(image_2)

        with torch.no_grad():
            # 获取预测结果
            output = self.net([photo_1, photo_2])[0]
            output = torch.nn.Sigmoid()(output)

        # 传入了保存图片的名字，才保存对比图
        if self.save_result_img_dir != "":
            self._save_comparison_result(image_1, image_2, output)

        # 手动释放显存
        if self.cuda:
            del photo_1, photo_2
            torch.cuda.empty_cache()

        return output

    def _save_comparison_result(self, image_1, image_2, similarity, show_result=False):
        """
        保存对比结果到文件，并可选是否显示结果。

        Args:
            image_1 (PIL.Image): 第一张图片。
            image_2 (PIL.Image): 第二张图片。
            similarity (torch.Tensor): 相似度值。
            show_result (bool): 是否显示结果（默认不显示）。
        """
        plt.figure(figsize=(10, 5))
        plt.subplot(1, 2, 1)
        plt.imshow(np.array(image_1))
        plt.title("Image 1")

        plt.subplot(1, 2, 2)
        plt.imshow(np.array(image_2))
        plt.title("Image 2")
        plt.text(-12, -12, f"Similarity: {similarity.item():.3f}", ha="center", va="bottom", fontsize=11)

        # 保存结果到文件
        save_path = f"{self.save_result_img_dir}.png"
        plt.savefig(save_path, bbox_inches="tight", dpi=300)
        logging.info(f"结果已保存到: {save_path}")

        # 根据参数决定是否显示结果
        if show_result:
            plt.show()
        plt.clf()  # 清理当前图像
        plt.close('all')  # 关闭所有图像，避免内存泄漏

    def convert_to_rgb(self, image):
        """将图像转换为 RGB 格式"""
        if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
            return image
        else:
            return image.convert("RGB")

    def preprocess_input(self, image):
        """归一化图像"""
        return image / 255.0


class SiameseImageComparator:
    """
    图像相似度比较器，支持多种比较场景。
    """

    def __init__(self, model_path, match_model, save_img_name="save_img_results"):
        """
        初始化比较器。

        Args:
            model_path (str): 模型权重路径。
            match_model (torch.nn.Module): 模型实例。
            save_img_name (str): 保存结果的文件名。
        """
        self.model = SiamesePredictor(model_path=model_path, match_model=match_model, save_result_img_name=save_img_name)

    def compare_images(self, image_path_1: str, image_path_2: str) -> float:
        """比较两张图片的相似度。

        Args:
            image_path_1 (str): 第一张图片的路径。
            image_path_2 (str): 第二张图片的路径。

        Returns:
            float: 相似度值。
        """
        image_1 = Image.open(image_path_1)
        image_2 = Image.open(image_path_2)
        similarity = self.model.detect_image(image_1, image_2)
        return float(f"{similarity.item():.6f}")